#include "AclProcess.h"

AclProcess::AclProcess()
{

}
AclProcess::~AclProcess()
{
    m_modelProcess = nullptr;
    aclError ret = aclrtSynchronizeStream(stream_);
    if (ret != ACL_ERROR_NONE) {
        cout << "some tasks in stream not done, ret = " << ret <<endl;
    }
    cout << "all tasks in stream done" << endl;
    ret = aclrtDestroyStream(stream_);
    if (ret != ACL_ERROR_NONE) {
        cout << "Destroy Stream faild, ret = " << ret <<endl;
    }
    cout << "Destroy Stream successfully" << endl;
    ret = aclrtDestroyContext(context_);
    if (ret != ACL_ERROR_NONE) {
        cout << "Destroy Context faild, ret = " << ret <<endl;
    }
    cout << "Destroy Context successfully" << endl;
    ret = aclFinalize();
    if (ret != ACL_ERROR_NONE) {
        cout << "Failed to deinit acl, ret = " << ret <<endl;
    }
    cout << "acl deinit successfully" << endl;
}

int AclProcess::Init(int deviceId,string cfgPath, string modelPath)
{
    //Init
    aclError ret = aclInit(cfgPath.c_str()); // Initialize ACL
    if (ret != ACL_ERROR_NONE) {
        cout << "Failed to init acl, ret = " << ret <<endl;
        return ret;
    }
    cout << "acl init successfully" << endl;
    ret = aclrtCreateContext(&context_, deviceId);
    if (ret != ACL_ERROR_NONE) {
        cout << "Failed to set current context, ret = " << ret << endl;
        return ret;
    }
    cout << "Create context successfully" << endl;
    ret = aclrtSetCurrentContext(context_);
    if (ret != ACL_ERROR_NONE) {
        cout << "Failed to set current context, ret = " << ret << endl;
        return ret;
    }
    cout << "set context successfully" << endl;
    ret = aclrtCreateStream(&stream_);
    if (ret != ACL_ERROR_NONE) {
        cout << "Failed to create stream, ret = " << ret << endl;
        return ret;
    }
    cout << "Create stream successfully" << endl;
    //Load model
    if (m_modelProcess == nullptr) {
        m_modelProcess = std::make_shared<ModelProcess>(deviceId, "");
    }
    ret = m_modelProcess->Init(modelPath);

    if (ret != ACL_ERROR_NONE) {
        cout << "Failed to initialize m_modelProcess, ret = " << ret << endl;
        return ret;
    }
    m_modelDesc = m_modelProcess->GetModelDesc();
    //get model input description and malloc them
    size_t inputSize = aclmdlGetNumInputs(m_modelDesc);
    for (size_t i = 0; i < inputSize; i++) {
        size_t bufferSize = aclmdlGetInputSizeByIndex(m_modelDesc, i);
        void *inputBuffer = nullptr;
        aclError ret = aclrtMalloc(&inputBuffer, bufferSize, ACL_MEM_MALLOC_NORMAL_ONLY);
        if (ret != ACL_ERROR_NONE) {
            cout << "Failed to malloc buffer, ret = " << ret << endl;
            return ret;
        }
        inputBuffers.push_back(inputBuffer);
        inputSizes.push_back(bufferSize);
    }
    //get model output description and malloc them
    size_t outputSize = aclmdlGetNumOutputs(m_modelDesc);
    for (size_t i = 0; i < outputSize; i++) {
        size_t bufferSize = aclmdlGetOutputSizeByIndex(m_modelDesc, i);
        void *outputBuffer = nullptr;
        aclError ret = aclrtMalloc(&outputBuffer, bufferSize, ACL_MEM_MALLOC_NORMAL_ONLY);
        if (ret != ACL_ERROR_NONE) {
            cout << "Failed to malloc buffer, ret = " << ret << endl;
            return ret;
        }
        outputBuffers.push_back(outputBuffer);
        outputSizes.push_back(bufferSize);
    }
    cout << "finish init AclProcess" << endl;
    return ACL_ERROR_NONE;
}

int AclProcess::Process(Mat& img)
{
    aclError ret = ACL_ERROR_NONE;
    Mat imgResize;
    //get input dims
    int origin_height,origin_width;
    origin_height = img.rows;
    origin_width = img.cols;
    int batch,channels,height,width;
    aclmdlIODims dims;
    aclmdlGetInputDims(m_modelDesc,0,&dims);
    if(dims.dimCount == 4){
        batch = dims.dims[0];
        height = dims.dims[1];
        width = dims.dims[2];
        channels = dims.dims[3];
    }
    /*warp affine image
    **what does the code do:               416
    **           1920                |-------------|
    **       _____________           |_____________|
    ** 1080 |  /______\   |  =====>  |  /______\   |
    **      |  |# n # |   |      416 |  |# n # |   |
    **      |-------------|          |-------------|
    **                               |_____________|
    **if your input image is with a static size
    **you can accelerate this process with aipp->padding
    **more detail:https://support.huaweicloud.com/ti-atc-A300_3000_3010/altasatc_16_009.html*/
    Mat M;
    float r = origin_height * 1.0 / origin_width;
    if(r > 1)
    {
        float offset = (width - origin_width * (height * 1.0 / origin_height)) / 2.0;
        M = getAffineTransform(
                    vector<Point2f>({Point2f(0, 0),Point2f(0, origin_height),Point2f(origin_width, 0)}),
                    vector<Point2f>({Point2f(offset, 0),Point2f(offset, height),Point2f((width - offset), 0)})
                    );
    }
    else
    {
        float offset = (height - origin_height * (width * 1.0 / origin_width)) / 2.0;
        M = getAffineTransform(
                    vector<Point2f>({Point2f(0, 0),Point2f(origin_width, 0),Point2f(0, origin_height)}),
                    vector<Point2f>({Point2f(0, offset),Point2f(width, offset),Point2f(0, (height-offset))})
                    );
    }
    warpAffine(img,imgResize,M,Size(height,width),INTER_NEAREST);
    //warp image infomation data
    float imgInfo[4] = {0};
    imgInfo[0] = height * 1.0;
    imgInfo[1] = width * 1.0;
    imgInfo[2] = origin_height * 1.0;
    imgInfo[3] = origin_width * 1.0;
    aclrtMemcpy(inputBuffers[0], inputSizes[0], imgResize.data,imgResize.cols * imgResize.rows * imgResize.channels(), ACL_MEMCPY_HOST_TO_DEVICE);
    aclrtMemcpy(inputBuffers[1], inputSizes[1], imgInfo,sizeof(float) * 4, ACL_MEMCPY_HOST_TO_DEVICE);

    //forward
    ret = m_modelProcess->ModelInference(inputBuffers, inputSizes, outputBuffers, outputSizes);
    if (ret != ACL_ERROR_NONE) {
        cout<<"model run faild.ret = "<< ret <<endl;
        return ret;
    }
    //postprocess
    cout << "begin postprocess"<<endl;
    YoloV3PostProcess(img, outputBuffers, outputSizes);
    cout<<"model run success!"<<endl;
    return ACL_ERROR_NONE;
}

aclError AclProcess::YoloV3PostProcess(Mat& img, std::vector<void *> outputBuffers, std::vector<size_t> outputSizes)
{
    // Get output of yolov3 model inference
    // According to the Operator Specification, YoloV3DetectionOutput as follows:
    // boxout: An NCHW tensor of type float16, describing the information of each output box,
    //         including the coordinates, class, and confidence.
    // boxoutnum: An NCHW tensor of type int32, specifying the number of output boxes.
    const size_t outputLen = outputSizes.size();
    if (outputLen <= 0) {
        cout << "Failed to get model output data" << endl;;
        return -4;
    }
    cout << "The number of output buffers of yolov3 model is " << outputLen << endl;

    aclError ret;
    void* hostPtr[outputLen];
    for (size_t j = 0; j < outputLen; j++) {
        aclError ret = (aclError)aclrtMallocHost(&hostPtr[j], outputSizes[j]);
        if (ret != ACL_ERROR_NONE) {
            cout << "Failed to malloc output buffer of model on host, ret = " << ret << endl;
            return ret;
        }
        ret = (aclError)aclrtMemcpy(hostPtr[j], outputSizes[j], outputBuffers[j],
                                     outputSizes[j], ACL_MEMCPY_DEVICE_TO_HOST);
        if (ret != ACL_ERROR_NONE) {
            cout << "Failed to copy output buffer of model from device to host, ret = " << ret << endl;
            return ret;
        }
    }

    // The number of detected targets is the number of output boxes, which is stored in hostPtr[1]
    uint32_t objNum = ((uint32_t*)hostPtr[1])[0];
    //uint32_t objNum = (uint32_t)((float*)hostPtr[1])[0];
    cout << "Object detected number is " << objNum << endl;
    ObjDetectInfo objInfo[objNum];
    // The coordinates, confidence and category of each target are stored in the memory pointed to by hostPtr[0] in turn
    for (uint32_t k = 0; k < objNum; k++) {
        int pos = 0;
        objInfo[k].leftTopX = ((float *)hostPtr[0])[objNum * (pos++) + k];
        objInfo[k].leftTopY = ((float *)hostPtr[0])[objNum * (pos++) + k];
        objInfo[k].rightBotX = ((float *)hostPtr[0])[objNum * (pos++) + k];
        objInfo[k].rightBotY = ((float *)hostPtr[0])[objNum * (pos++) + k];
        objInfo[k].confidence = ((float *)hostPtr[0])[objNum * (pos++) + k];
        objInfo[k].classId = ((float *)hostPtr[0])[objNum * (pos++) + k];
        cout << "#Obj" << k << ", " << "box(" << objInfo[k].leftTopX << ", " << objInfo[k].leftTopY << ", "
                << objInfo[k].rightBotX << ", " << objInfo[k].rightBotY << ")  "
                << " confidence: " << objInfo[k].confidence << " lable: " << objInfo[k].classId << endl;
        rectangle(img, Point(objInfo[k].leftTopX, objInfo[k].leftTopY), Point(objInfo[k].rightBotX, objInfo[k].rightBotY), Scalar(0,255,0), 2);
    }

    return ACL_ERROR_NONE;
}
